Qdrant: ベクトル検索入門チュートリアル

#pip install qdrant-client sentence-transformers python-dotenv -q

import os
from dotenv import load_dotenv
from qdrant_client import QdrantClient, models
from sentence_transformers import SentenceTransformer

# Load environment variables from .env file
load_dotenv()

# Get the API key and URL from environment variables
QDRANT_API_KEY = os.getenv("QDRANT_API_KEY")
QDRANT_URL = os.getenv("QDRANT_URL")

# Initialize the Qdrant client for in-memory storage
# client = QdrantClient(":memory:")

# If you were connecting to Qdrant Cloud, you would use:
client = QdrantClient(url=QDRANT_URL, api_key=QDRANT_API_KEY)

my_collection = "my_first_collection"

client.recreate_collection(
    collection_name=my_collection,
    vectors_config=models.VectorParams(size=384, distance=models.Distance.COSINE),
)

# Load a pre-trained sentence transformer model
model = SentenceTransformer('all-MiniLM-L6-v2')

# Let's create some sample documents
documents = [
    {"id": 1, "text": "Qdrant is a vector database for building AI applications.", "metadata": {"type": "tech"}},
    {"id": 2, "text": "The Eiffel Tower is a famous landmark in Paris, France.", "metadata": {"type": "travel"}},
    {"id": 3, "text": "A vector database indexes vectors for easy search and retrieval.", "metadata": {"type": "tech"}},
    {"id": 4, "text": "The Great Wall of China is one of the world's wonders.", "metadata": {"type": "travel"}},
    {"id": 5, "text": "Artificial intelligence is transforming many industries.", "metadata": {"type": "tech"}}
]

# Generate embeddings for our documents
embeddings = model.encode([doc["text"] for doc in documents])

client.upsert(
    collection_name=my_collection,
    points=[
        models.PointStruct(
            id=doc["id"],
            vector=embedding.tolist(),
            payload=doc["metadata"]
        )
        for doc, embedding in zip(documents, embeddings)
    ],
    wait=True,
)

query = "What is a vector database?"
query_vector = model.encode(query).tolist()

search_results = client.query_points(
    collection_name=my_collection,
    query=query_vector,
    limit=2,  # Return the top 2 most similar results
    with_payload=True,
)

# Map IDs to original document text so we can display it with results
id_to_text = {doc["id"]: doc["text"] for doc in documents}

print("Semantic Search Results:")
for result in search_results.points:
    print(f"- ID: {result.id}, Score: {result.score:.4f}")
    text = id_to_text.get(result.id)
    if text is not None:
        print(f"  Text: {text}")

# Create a payload index for the string field 'type' so we can filter on it
# Qdrant requires an index of type KEYWORD for exact-match string filters
try:
    client.create_payload_index(
        collection_name=my_collection,
        field_name="type",
        field_schema=models.PayloadSchemaType.KEYWORD,
    )
    print("Created payload index for 'type' (KEYWORD).")
except Exception as e:
    # Safe to ignore if already exists and we're re-running cells
    if "already exists" in str(e).lower():
        print("Payload index for 'type' already exists.")
    else:
        raise

filtered_search_results = client.query_points(
    collection_name=my_collection,
    query=query_vector,
    query_filter=models.Filter(
        must=[
            models.FieldCondition(
                key="type",
                match=models.MatchValue(value="travel")
            )
        ]
    ),
    limit=2,
    with_payload=True,
)

print("Filtered Search Results (type='travel'):")
for result in filtered_search_results.points:
    print(f"- ID: {result.id}, Score: {result.score:.4f}, Payload: {result.payload}")
    text = id_to_text.get(result.id)
    if text is not None:
        print(f"  Text: {text}")